"""
动态密钥轮换模块 - 定期更新签名密钥

这个模块实现了动态密钥轮换功能：
- 自动密钥生成和轮换
- 密钥版本管理
- 无缝密钥切换
- 密钥分发和同步
- 密钥撤销和清理
"""

import asyncio
import logging
import time
import secrets
import hashlib
from dataclasses import dataclass, field
from datetime import datetime, timezone, timedelta
from enum import Enum
from typing import Any, Dict, List, Optional, Union, Callable
import json
import base64
from pathlib import Path

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa, ed25519
from cryptography.hazmat.primitives.kdf import pbkdf2
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

logger = logging.getLogger(__name__)


class KeyType(Enum):
    """密钥类型"""
    RSA = "rsa"
    ED25519 = "ed25519"
    AES = "aes"
    HMAC = "hmac"


class KeyStatus(Enum):
    """密钥状态"""
    ACTIVE = "active"        # 当前使用
    PENDING = "pending"      # 待激活
    EXPIRED = "expired"      # 已过期
    REVOKED = "revoked"      # 已撤销
    ARCHIVED = "archived"    # 已归档


class RotationStrategy(Enum):
    """轮换策略"""
    TIME_BASED = "time_based"      # 基于时间
    USAGE_BASED = "usage_based"    # 基于使用量
    MANUAL = "manual"              # 手动触发
    EVENT_BASED = "event_based"    # 基于事件


@dataclass
class KeyInfo:
    """密钥信息"""
    key_id: str
    key_type: KeyType
    status: KeyStatus
    created_at: datetime
    expires_at: datetime
    last_used: Optional[datetime] = None
    usage_count: int = 0
    version: int = 1
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class KeyRotationConfig:
    """密钥轮换配置"""
    # 轮换策略
    rotation_strategy: RotationStrategy = RotationStrategy.TIME_BASED
    rotation_interval: int = 86400  # 24小时（秒）
    key_lifetime: int = 7 * 86400   # 7天（秒）
    overlap_period: int = 3600      # 1小时重叠期（秒）
    
    # 密钥配置
    key_types: List[KeyType] = field(default_factory=lambda: [KeyType.RSA, KeyType.HMAC])
    rsa_key_size: int = 2048
    aes_key_size: int = 256
    hmac_key_size: int = 256
    
    # 存储配置
    key_storage_path: str = "keys"
    max_key_versions: int = 5
    enable_key_encryption: bool = True
    key_encryption_password: Optional[str] = None
    
    # 分发配置
    enable_key_distribution: bool = True
    distribution_timeout: float = 30.0
    max_distribution_retries: int = 3
    
    # 监控配置
    enable_usage_tracking: bool = True
    usage_threshold: int = 10000  # 使用次数阈值
    enable_audit_log: bool = True


class KeyManager:
    """密钥管理器"""
    
    def __init__(self, config: Optional[KeyRotationConfig] = None):
        self.config = config or KeyRotationConfig()
        
        # 密钥存储
        self.keys: Dict[str, Dict[str, Any]] = {}  # key_id -> {info, key_data}
        self.active_keys: Dict[KeyType, str] = {}  # key_type -> active_key_id
        self.pending_keys: Dict[KeyType, str] = {}  # key_type -> pending_key_id
        
        # 密钥使用统计
        self.key_usage_stats: Dict[str, int] = {}
        
        # 轮换任务
        self.rotation_task: Optional[asyncio.Task] = None
        self.is_running = False
        
        # 锁
        self.lock = asyncio.Lock()
        
        # 密钥存储目录
        self.key_dir = Path(self.config.key_storage_path)
        self.key_dir.mkdir(exist_ok=True)
        
        # 回调函数
        self.key_rotation_callbacks: List[Callable] = []
        self.key_distribution_callbacks: List[Callable] = []
        
        self.logger = logging.getLogger(__name__)
        self.logger.info("密钥管理器已初始化")
    
    async def start(self) -> None:
        """启动密钥管理器"""
        if self.is_running:
            return
        
        self.is_running = True
        
        # 加载现有密钥
        await self._load_existing_keys()
        
        # 启动轮换任务
        if self.config.rotation_strategy != RotationStrategy.MANUAL:
            self.rotation_task = asyncio.create_task(self._rotation_worker())
        
        self.logger.info("密钥管理器已启动")
    
    async def stop(self) -> None:
        """停止密钥管理器"""
        if not self.is_running:
            return
        
        self.is_running = False
        
        if self.rotation_task:
            self.rotation_task.cancel()
            try:
                await self.rotation_task
            except asyncio.CancelledError:
                pass
        
        self.logger.info("密钥管理器已停止")
    
    async def generate_key(self, 
                          key_type: KeyType,
                          key_id: Optional[str] = None) -> str:
        """生成新密钥"""
        if key_id is None:
            key_id = self._generate_key_id(key_type)
        
        try:
            key_data = await self._create_key(key_type)
            
            # 创建密钥信息
            key_info = KeyInfo(
                key_id=key_id,
                key_type=key_type,
                status=KeyStatus.PENDING,
                created_at=datetime.now(timezone.utc),
                expires_at=datetime.now(timezone.utc) + timedelta(seconds=self.config.key_lifetime),
                version=1
            )
            
            # 存储密钥
            async with self.lock:
                self.keys[key_id] = {
                    "info": key_info,
                    "data": key_data
                }
                
                # 保存到文件
                await self._save_key_to_file(key_id, key_info, key_data)
            
            self.logger.info(f"密钥生成成功: {key_id} ({key_type.value})")
            return key_id
            
        except Exception as e:
            self.logger.error(f"密钥生成失败: {e}")
            raise
    
    async def _create_key(self, key_type: KeyType) -> bytes:
        """创建密钥数据"""
        if key_type == KeyType.RSA:
            private_key = rsa.generate_private_key(
                public_exponent=65537,
                key_size=self.config.rsa_key_size,
                backend=default_backend()
            )
            return private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            )
        
        elif key_type == KeyType.ED25519:
            private_key = ed25519.Ed25519PrivateKey.generate()
            return private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            )
        
        elif key_type == KeyType.AES:
            return secrets.token_bytes(self.config.aes_key_size // 8)
        
        elif key_type == KeyType.HMAC:
            return secrets.token_bytes(self.config.hmac_key_size // 8)
        
        else:
            raise ValueError(f"不支持的密钥类型: {key_type}")
    
    async def activate_key(self, key_id: str) -> bool:
        """激活密钥"""
        async with self.lock:
            if key_id not in self.keys:
                return False
            
            key_entry = self.keys[key_id]
            key_info = key_entry["info"]
            
            # 检查密钥是否有效
            if key_info.status not in [KeyStatus.PENDING, KeyStatus.ACTIVE]:
                return False
            
            # 设置新密钥为活跃状态
            key_info.status = KeyStatus.ACTIVE
            key_info.last_used = datetime.now(timezone.utc)
            
            # 更新活跃密钥映射
            self.active_keys[key_info.key_type] = key_id
            
            # 如果有旧密钥，设置为过期
            old_key_id = self.pending_keys.get(key_info.key_type)
            if old_key_id and old_key_id != key_id:
                old_key_info = self.keys[old_key_id]["info"]
                old_key_info.status = KeyStatus.EXPIRED
            
            self.pending_keys[key_info.key_type] = key_id
            
            # 保存更改
            await self._save_key_to_file(key_id, key_info, key_entry["data"])
            
            self.logger.info(f"密钥已激活: {key_id}")
            return True
    
    async def revoke_key(self, key_id: str) -> bool:
        """撤销密钥"""
        async with self.lock:
            if key_id not in self.keys:
                return False
            
            key_entry = self.keys[key_id]
            key_info = key_entry["info"]
            
            # 设置状态为已撤销
            key_info.status = KeyStatus.REVOKED
            
            # 如果这是活跃密钥，需要激活下一个密钥
            if self.active_keys.get(key_info.key_type) == key_id:
                # 查找下一个可用的密钥
                next_key_id = await self._find_next_available_key(key_info.key_type)
                if next_key_id:
                    await self.activate_key(next_key_id)
                else:
                    # 生成新密钥
                    new_key_id = await self.generate_key(key_info.key_type)
                    await self.activate_key(new_key_id)
            
            # 保存更改
            await self._save_key_to_file(key_id, key_info, key_entry["data"])
            
            self.logger.info(f"密钥已撤销: {key_id}")
            return True
    
    async def get_active_key(self, key_type: KeyType) -> Optional[bytes]:
        """获取活跃密钥"""
        async with self.lock:
            active_key_id = self.active_keys.get(key_type)
            if not active_key_id or active_key_id not in self.keys:
                return None
            
            key_entry = self.keys[active_key_id]
            key_info = key_entry["info"]
            
            # 更新使用统计
            if self.config.enable_usage_tracking:
                key_info.usage_count += 1
                key_info.last_used = datetime.now(timezone.utc)
                self.key_usage_stats[active_key_id] = self.key_usage_stats.get(active_key_id, 0) + 1
            
            return key_entry["data"]
    
    async def get_key_info(self, key_id: str) -> Optional[KeyInfo]:
        """获取密钥信息"""
        async with self.lock:
            if key_id in self.keys:
                return self.keys[key_id]["info"]
            return None
    
    async def list_keys(self, 
                       key_type: Optional[KeyType] = None,
                       status: Optional[KeyStatus] = None) -> List[KeyInfo]:
        """列出密钥"""
        async with self.lock:
            keys = []
            for key_entry in self.keys.values():
                key_info = key_entry["info"]
                
                if key_type and key_info.key_type != key_type:
                    continue
                
                if status and key_info.status != status:
                    continue
                
                keys.append(key_info)
            
            return keys
    
    async def _rotation_worker(self) -> None:
        """密钥轮换工作线程"""
        while self.is_running:
            try:
                await self._check_and_rotate_keys()
                await asyncio.sleep(self.config.rotation_interval)
            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"密钥轮换异常: {e}")
                await asyncio.sleep(60)
    
    async def _check_and_rotate_keys(self) -> None:
        """检查并轮换密钥"""
        current_time = datetime.now(timezone.utc)
        
        for key_type in self.config.key_types:
            try:
                # 检查是否需要轮换
                if await self._should_rotate_key(key_type, current_time):
                    await self._rotate_key(key_type)
            except Exception as e:
                self.logger.error(f"密钥轮换失败: {key_type.value}, 错误: {e}")
    
    async def _should_rotate_key(self, key_type: KeyType, current_time: datetime) -> bool:
        """检查是否应该轮换密钥"""
        active_key_id = self.active_keys.get(key_type)
        if not active_key_id:
            return True  # 没有活跃密钥，需要生成
        
        key_entry = self.keys[active_key_id]
        key_info = key_entry["info"]
        
        # 检查过期时间
        if current_time >= key_info.expires_at - timedelta(seconds=self.config.overlap_period):
            return True
        
        # 检查使用量
        if (self.config.rotation_strategy == RotationStrategy.USAGE_BASED and
            key_info.usage_count >= self.config.usage_threshold):
            return True
        
        return False
    
    async def _rotate_key(self, key_type: KeyType) -> None:
        """轮换密钥"""
        try:
            # 生成新密钥
            new_key_id = await self.generate_key(key_type)
            
            # 激活新密钥
            await self.activate_key(new_key_id)
            
            # 触发轮换回调
            await self._trigger_rotation_callbacks(key_type, new_key_id)
            
            # 分发新密钥
            if self.config.enable_key_distribution:
                await self._distribute_key(new_key_id)
            
            self.logger.info(f"密钥轮换完成: {key_type.value} -> {new_key_id}")
            
        except Exception as e:
            self.logger.error(f"密钥轮换失败: {key_type.value}, 错误: {e}")
            raise
    
    async def _find_next_available_key(self, key_type: KeyType) -> Optional[str]:
        """查找下一个可用的密钥"""
        for key_id, key_entry in self.keys.items():
            key_info = key_entry["info"]
            if (key_info.key_type == key_type and 
                key_info.status == KeyStatus.PENDING):
                return key_id
        return None
    
    async def _load_existing_keys(self) -> None:
        """加载现有密钥"""
        try:
            for key_file in self.key_dir.glob("*.json"):
                try:
                    with open(key_file, 'r') as f:
                        key_data = json.load(f)
                    
                    key_id = key_data["info"]["key_id"]
                    key_info = KeyInfo(**key_data["info"])
                    
                    # 加载密钥数据
                    key_data_file = self.key_dir / f"{key_id}.key"
                    if key_data_file.exists():
                        with open(key_data_file, 'rb') as f:
                            encrypted_data = f.read()
                        
                        # 解密密钥数据
                        if self.config.enable_key_encryption:
                            key_data_bytes = await self._decrypt_key_data(encrypted_data)
                        else:
                            key_data_bytes = encrypted_data
                        
                        self.keys[key_id] = {
                            "info": key_info,
                            "data": key_data_bytes
                        }
                        
                        # 更新活跃密钥映射
                        if key_info.status == KeyStatus.ACTIVE:
                            self.active_keys[key_info.key_type] = key_id
                        
                        self.logger.info(f"加载密钥: {key_id}")
                
                except Exception as e:
                    self.logger.error(f"加载密钥文件失败: {key_file}, 错误: {e}")
        
        except Exception as e:
            self.logger.error(f"加载现有密钥失败: {e}")
    
    async def _save_key_to_file(self, key_id: str, key_info: KeyInfo, key_data: bytes) -> None:
        """保存密钥到文件"""
        try:
            # 保存密钥信息
            info_file = self.key_dir / f"{key_id}.json"
            with open(info_file, 'w') as f:
                json.dump({
                    "info": {
                        "key_id": key_info.key_id,
                        "key_type": key_info.key_type.value,
                        "status": key_info.status.value,
                        "created_at": key_info.created_at.isoformat(),
                        "expires_at": key_info.expires_at.isoformat(),
                        "last_used": key_info.last_used.isoformat() if key_info.last_used else None,
                        "usage_count": key_info.usage_count,
                        "version": key_info.version,
                        "metadata": key_info.metadata
                    }
                }, f, indent=2)
            
            # 保存密钥数据
            data_file = self.key_dir / f"{key_id}.key"
            if self.config.enable_key_encryption:
                encrypted_data = await self._encrypt_key_data(key_data)
            else:
                encrypted_data = key_data
            
            with open(data_file, 'wb') as f:
                f.write(encrypted_data)
        
        except Exception as e:
            self.logger.error(f"保存密钥文件失败: {key_id}, 错误: {e}")
            raise
    
    async def _encrypt_key_data(self, key_data: bytes) -> bytes:
        """加密密钥数据"""
        if not self.config.key_encryption_password:
            return key_data
        
        # 生成随机IV
        iv = secrets.token_bytes(16)
        
        # 生成密钥
        password = self.config.key_encryption_password.encode()
        salt = secrets.token_bytes(16)
        kdf = pbkdf2.PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
            backend=default_backend()
        )
        key = kdf.derive(password)
        
        # 加密数据
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
        encryptor = cipher.encryptor()
        
        # 填充数据
        padding_length = 16 - (len(key_data) % 16)
        padded_data = key_data + bytes([padding_length] * padding_length)
        
        encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
        
        # 返回 salt + iv + encrypted_data
        return salt + iv + encrypted_data
    
    async def _decrypt_key_data(self, encrypted_data: bytes) -> bytes:
        """解密密钥数据"""
        if not self.config.key_encryption_password:
            return encrypted_data
        
        # 提取 salt, iv 和加密数据
        salt = encrypted_data[:16]
        iv = encrypted_data[16:32]
        encrypted = encrypted_data[32:]
        
        # 生成密钥
        password = self.config.key_encryption_password.encode()
        kdf = pbkdf2.PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
            backend=default_backend()
        )
        key = kdf.derive(password)
        
        # 解密数据
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
        decryptor = cipher.decryptor()
        
        decrypted_data = decryptor.update(encrypted) + decryptor.finalize()
        
        # 移除填充
        padding_length = decrypted_data[-1]
        return decrypted_data[:-padding_length]
    
    def _generate_key_id(self, key_type: KeyType) -> str:
        """生成密钥ID"""
        timestamp = int(time.time())
        random_part = secrets.token_hex(8)
        return f"{key_type.value}_{timestamp}_{random_part}"
    
    async def _trigger_rotation_callbacks(self, key_type: KeyType, new_key_id: str) -> None:
        """触发轮换回调"""
        for callback in self.key_rotation_callbacks:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(key_type, new_key_id)
                else:
                    callback(key_type, new_key_id)
            except Exception as e:
                self.logger.error(f"轮换回调执行失败: {e}")
    
    async def _distribute_key(self, key_id: str) -> None:
        """分发密钥"""
        for callback in self.key_distribution_callbacks:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(key_id)
                else:
                    callback(key_id)
            except Exception as e:
                self.logger.error(f"密钥分发回调执行失败: {e}")
    
    def add_rotation_callback(self, callback: Callable) -> None:
        """添加轮换回调"""
        self.key_rotation_callbacks.append(callback)
    
    def add_distribution_callback(self, callback: Callable) -> None:
        """添加分发回调"""
        self.key_distribution_callbacks.append(callback)
    
    async def get_rotation_status(self) -> Dict[str, Any]:
        """获取轮换状态"""
        status = {
            "is_running": self.is_running,
            "active_keys": {},
            "pending_keys": {},
            "key_counts": {},
            "usage_stats": self.key_usage_stats
        }
        
        # 活跃密钥
        for key_type, key_id in self.active_keys.items():
            key_info = self.keys[key_id]["info"]
            status["active_keys"][key_type.value] = {
                "key_id": key_id,
                "created_at": key_info.created_at.isoformat(),
                "expires_at": key_info.expires_at.isoformat(),
                "usage_count": key_info.usage_count,
                "last_used": key_info.last_used.isoformat() if key_info.last_used else None
            }
        
        # 待激活密钥
        for key_type, key_id in self.pending_keys.items():
            if key_id not in self.active_keys.get(key_type, ""):
                key_info = self.keys[key_id]["info"]
                status["pending_keys"][key_type.value] = {
                    "key_id": key_id,
                    "created_at": key_info.created_at.isoformat(),
                    "expires_at": key_info.expires_at.isoformat()
                }
        
        # 密钥计数
        for key_type in KeyType:
            count = sum(1 for key_entry in self.keys.values() 
                       if key_entry["info"].key_type == key_type)
            status["key_counts"][key_type.value] = count
        
        return status


# 全局密钥管理器
_global_key_manager: Optional[KeyManager] = None


def get_key_manager() -> KeyManager:
    """获取全局密钥管理器"""
    global _global_key_manager
    if _global_key_manager is None:
        _global_key_manager = KeyManager()
    return _global_key_manager
